#  -*- coding: utf-8 -*-

import traceback
import matplotlib.pyplot as plt
import numpy as np
from pymongo import UpdateOne
import pandas as pd
from data.data_module import DataModule
from trading.signal.computer.base_signal_computer import BaseSignalComputer
from util.stock_util import get_all_codes
import mpl_finance as mpf


class BollSignalComputer(BaseSignalComputer):
    def __init__(self):
        BaseSignalComputer.__init__(self, 'boll_signal')

    def compute(self, begin_date, end_date):
        """
        计算指定日期内的信号
        :param begin_date: 开始日期
        :param end_date: 结束日期
        """
        # all_codes = get_all_codes('2018-08-24')
        all_codes = ['601398']
        dm = DataModule()

        N = 20
        k = 2
        mistake = 0.1
        for code in all_codes[:]:
            print(code)
            try:
                df_daily = dm.get_k_data(code, begin_date=begin_date, autype='kba', end_date=end_date, trade=True)
                lowlist = df_daily['low'].rolling(9).min()
                highlist = df_daily['high'].rolling(9).max()
                def ee(x):
                    if x>100:
                        return 100
                    else:
                        return x
                rsv =((df_daily['close'].astype('float')-lowlist)/(highlist-lowlist))*100
                df_daily['k']=(pd.ewma(rsv,com=2)).apply(ee)
                df_daily['d']=(pd.ewma( df_daily['k'],com=2)).apply(ee)
                df_daily['j']=(3*df_daily['k']-2*df_daily['d']).apply(ee)
                print (df_daily)
                plt.plot(df_daily['k'])
                plt.plot(df_daily['d'])
                plt.plot(df_daily['j'])
                plt.xticks([])
                plt.show()
            except:
                traceback.print_exc()
                # break


if __name__ == '__main__':
    BollSignalComputer().compute(begin_date='2018-01-01', end_date='2018-08-30')
